# 0. 引入必要的包  
import os  
import numpy as np  
from sklearn.model_selection import train_test_split  
from util import get, preprocess_image  
from joblib import dump  
from tqdm import tqdm  
import time  
import glob  
  
# 1. 读取配置文件中的信息
train_dir = get("train") # 获取 训练数据路径
char_styles = get("char_styles") # 获取 字符样式列表，注意: 必须是列标
new_size = get("new_size") # 获取 新图像大小元组, 注意: 必须包含h和w

# 2. 生成X,y   
print("# 读取训练数据并进行预处理，")  
X_images, y_labels = [], []  
# 循环遍历每个书法体类别  
for style in char_styles:  
    # 获取当前类别的所有图片路径  
    image_files = glob.glob(os.path.join(train_dir, f"train_{style}*"))  
    # 循环遍历每个图片  
    for file_path in tqdm(image_files, desc=f"处理 {style} 图像"):  
        # 调用preprocess_image函数处理图像  
        img = preprocess_image(file_path, new_size)  
        # 把预处理过的图像添加到X中  
        X_images.append(img)  
        # 获取图像对应的类别添加到y中，类别索引与 char_styles 列表中的顺序相对应
        y_labels.append(char_styles.index(style))  
  
# 将数据转换为NumPy数组以便于后续处理  
X = np.array(X_images)  
y = np.array(y_labels, dtype=np.int64)  
  
# 3. 分割测试集和训练集  
print("# 将数据按 80% 和 20% 的比例分割")  
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)  
  
# 4. 打印样本维度和类型信息
print("X_train: ", X_train.shape, X_train.dtype)  # 训练集特征的维度和类型
print("X_test: ", X_test.shape, X_test.dtype)  # 测试集特征的维度和类型
print("y_train: ", y_train.shape, y_train.dtype)  # 训练集标签的维度和类型
print("y_test: ", y_test.shape, y_test.dtype)  # 测试集标签的维度和类型

  
# 5. 序列化分割后的训练和测试样本  
start_time = time.time()   
output_dir = './Xys'  
os.makedirs(output_dir, exist_ok=True)  
file_path = os.path.join(output_dir, 'Xy.joblib')  
Xy = (X_train, X_test, y_train, y_test)  
dump(Xy, file_path)  
  
# 打印文件位置和大小  
print(f"把(X_train, X_test, y_train, y_test)保存到 {file_path}")  
print(f"保存完毕, 文件位置: {file_path}, 大小: {os.path.getsize(file_path) / (1024 * 1024):.3f}M")  
  
# 计算并打印运行时间  
end_time = time.time()  
print(f"运行时间: {end_time - start_time:.3f}秒")